#include "utils.h"
#include <fstream>
#include <algorithm>
#include <iostream>

using namespace mindspore;

std::vector<std::string> GetAllBins(std::string_view dir_name) {
  struct dirent *filename;  // return value for readdir()
  DIR *dir = OpenDir(dir_name);                // return value for opendir()
  if (dir == nullptr) {
    return {};
  }

  /* read all the files in the dir ~ */
  std::vector<std::string> res;
  while ((filename = readdir(dir)) != nullptr) {
    std::string d_name = std::string(filename->d_name);
    // get rid of "." and ".."
    if (d_name == "." ||
        d_name == ".." ||
        d_name.size() <= 3 ||
        d_name.substr(d_name.size() - 4) != ".bin" ||
        filename->d_type != DT_REG)
      continue;
    res.emplace_back(std::string(dir_name) + "/" + filename->d_name);
  }

  std::sort(res.begin(), res.end());
  for (auto &f : res) {
    std::cout << "bin file: " << f << std::endl;
  }
  return res;
}

std::vector<std::string> GetAllFiles(std::string_view dir_name) {
  struct dirent *filename;  // return value for readdir()
  DIR *dir = OpenDir(dir_name);                // return value for opendir()
  if (dir == nullptr) {
    return {};
  }

  /* read all the files in the dir ~ */
  std::vector<std::string> res;
  while ((filename = readdir(dir)) != nullptr) {
    std::string d_name = std::string(filename->d_name);
    // get rid of "." and ".."
    if (d_name == "." ||
        d_name == ".." ||
        filename->d_type != DT_REG)
      continue;
    res.emplace_back(std::string(dir_name) + "/" + filename->d_name);
  }

  std::sort(res.begin(), res.end());
  for (auto &f : res) {
    std::cout << "bin file: " << f << std::endl;
  }
  return res;
}

Buffer ReadFile(const std::string &file) {
  Buffer buffer;
  if (file.empty()) {
    std::cout << "Pointer file is nullptr" << std::endl;
    return buffer;
  }

  std::ifstream ifs(file);
  if (!ifs.good()) {
    std::cout << "File: " << file << " is not exist" << std::endl;
    return buffer;
  }

  if (!ifs.is_open()) {
    std::cout << "File: " << file << "open failed" << std::endl;
    return buffer;
  }

  ifs.seekg(0, std::ios::end);
  size_t size = ifs.tellg();
  buffer.ResizeData(size);
  if (buffer.DataSize() != size) {
    std::cout << "Malloc buf failed, file: " << file << std::endl;
    ifs.close();
    return buffer;
  }

  ifs.seekg(0, std::ios::beg);
  ifs.read(reinterpret_cast<char *>(buffer.MutableData()), size);
  ifs.close();

  return buffer;
}

mindspore::MSTensor ReadFileToTensor(const std::string &file) {
  if (file.empty()) {
    std::cout << "Pointer file is nullptr" << std::endl;
    return mindspore::MSTensor();
  }

  std::ifstream ifs(file);
  if (!ifs.good()) {
    std::cout << "File: " << file << " is not exist" << std::endl;
    return mindspore::MSTensor();
  }

  if (!ifs.is_open()) {
    std::cout << "File: " << file << "open failed" << std::endl;
    return mindspore::MSTensor();
  }

  ifs.seekg(0, std::ios::end);
  size_t size = ifs.tellg();
  mindspore::MSTensor buffer("file", DataType::kNumberTypeUInt8, {static_cast<int64_t>(size)}, nullptr, size);

  ifs.seekg(0, std::ios::beg);
  ifs.read(reinterpret_cast<char *>(buffer.MutableData()), size);
  ifs.close();

  return buffer;
}

void SaveFile(int idx, MSTensor buffer, const std::string &file_path) {
  std::string file_name = file_path.substr(file_path.find_last_of("/") + 1);
  std::string path = "../result_Files/output_" + std::to_string(idx) + "_in_" + file_name;
  FILE *output_file = fopen(path.c_str(), "wb");
  if (output_file == nullptr) {
    std::cout << "Write file " << path << " failed when fopen" << std::endl;
    return;
  }

  size_t wsize = fwrite(buffer.MutableData(), buffer.DataSize(), sizeof(uint8_t), output_file);
  if (wsize == 0) {
    std::cout << "Write file " << path << " failed when fwrite" << std::endl;
    return;
  }

  fclose(output_file);
  std::cout << "Save file " << path << " length " << buffer.DataSize() << " success." << std::endl;
}

std::vector<std::vector<std::string>> GetAllInputData(std::string_view dir_name) {
  std::vector<std::vector<std::string>> ret;
  auto tmp = GetAllBins(dir_name);
  if (!tmp.empty()) {
    std::cout << dir_name << " is bin dir.";
    ret.emplace_back(tmp);
    return ret;
  }

  DIR *dir = OpenDir(dir_name);
  if (dir == nullptr) {
    return {};
  }
  struct dirent *filename;  // return value for readdir()
  /* read all the files in the dir ~ */
  std::vector<std::string> sub_dirs;
  while ((filename = readdir(dir)) != nullptr) {
    std::string d_name = std::string(filename->d_name);
    // get rid of "." and ".."
    if (d_name == "." ||
        d_name == ".." ||
        d_name.empty() ||
        d_name[0] != '0') {
      continue;
    }

    std::string dir_path = RealPath(std::string(dir_name) + "/" + filename->d_name);
    struct stat s;
    lstat(dir_path.c_str(), &s);
    if (!S_ISDIR(s.st_mode)) {
      continue;
    }

    sub_dirs.emplace_back(dir_path);
  }
  std::sort(sub_dirs.begin(), sub_dirs.end());
  for (auto &d : sub_dirs) {
    std::cout << "sub dir: " << d << std::endl;
  }
  for (auto &d : sub_dirs) {
    ret.emplace_back(GetAllBins(d));
  }

  return ret;
}

DIR *OpenDir(std::string_view dir_name) {
  // check the parameter !
  if (dir_name.empty()) {
    std::cout << " dir_name is null ! " << std::endl;
    return nullptr;
  }

  std::string real_path = RealPath(dir_name);

  // check if dir_name is a valid dir
  struct stat s;
  lstat(real_path.c_str(), &s);
  if (!S_ISDIR(s.st_mode)) {
    std::cout << "dir_name is not a valid directory !" << std::endl;
    return nullptr;
  }

  DIR *dir;
  dir = opendir(real_path.c_str());
  if (dir == nullptr) {
    std::cout << "Can not open dir " << dir_name << std::endl;
    return nullptr;
  }
  std::cout << "Successfully opened the dir " << dir_name << std::endl;
  return dir;
}

std::string RealPath(std::string_view path) {
  char real_path_mem[PATH_MAX] = {0};
  char *real_path_ret = nullptr;
  real_path_ret = realpath(path.data(), real_path_mem);

  if (real_path_ret == nullptr) {
    std::cout << "File: " << path << " is not exist.";
    return "";
  }

  std::string real_path(real_path_mem);
  std::cout << path << " realpath is: " << real_path << std::endl;
  return real_path;
}